import numpy as np
import cv2
import os

save_image_dir = "mask_dir/image/"
save_mask_dir = "mask_dir/mask/"
target_src_dir = "../test/"

wndName = "Effect window"
maskName = "Mask window"
pre_pt_x, pre_pt_y = -1, -1
cur_pt_x, cur_pt_y = -1, -1
s_bMouseLButtonDown = False
s_bMouseRButtonDown = False
line_start = None
line_down_start = None


def fillHole(src):
    im_floodfill = src.copy()
    h, w = src.shape[:2]
    mask = np.zeros((h + 2, w + 2), np.uint8)
    cv2.floodFill(im_floodfill, mask, (0, 0), 255)
    im_floodfill_inv = cv2.bitwise_not(im_floodfill)
    im_out = src | im_floodfill_inv
    return im_out


def draw_mask_image(src,mask):
    dst = src.copy()
    H, W, _ = dst.shape
    coloring = np.zeros(shape=(H, W, 3), dtype=np.uint8)
    coloring[:, :, 2] = mask[:, :]
    dst = cv2.addWeighted(dst, 1, coloring, 0.6, 0.)
    return dst


def draw_mask(event, x, y, flags, param):
    global pre_pt_x, pre_pt_y, cur_pt_x, cur_pt_y
    global s_bMouseLButtonDown, s_bMouseRButtonDown
    global dstImage, copyImage, tmpImage
    global line_start, line_down_start

    cur_pt_x, cur_pt_y = x, y

    if event == cv2.EVENT_LBUTTONDOWN and not (flags & cv2.EVENT_FLAG_RBUTTON):
            copyImage = dstImage.copy()
            pre_pt_x, pre_pt_y = x,y
            cv2.circle(copyImage, (pre_pt_x, pre_pt_y), 3, (0, 0, 255), -1, 0)
            cv2.imshow(wndName, copyImage)

            dstImage = copyImage.copy()
            s_bMouseLButtonDown = True

    elif event == cv2.EVENT_LBUTTONDOWN and (flags & cv2.EVENT_FLAG_RBUTTON):
            cv2.circle(copyImage, (x, y), 3, (0, 0, 255), -1, 0)
            cv2.line(dstImage, (line_start[0], line_start[1]), (x, y), (255, 0, 0), 3)
            cv2.line(tmpImage, (line_start[0], line_start[1]), (x, y), (255, 0, 0), 3)

            line_start = (x, y)

    elif event == cv2.EVENT_RBUTTONDOWN:
            copyImage = dstImage.copy()
            cv2.circle(copyImage, (x, y), 3, (0, 0, 255), -1, 0)
            cv2.imshow(wndName, copyImage)

            dstImage = copyImage.copy()
            s_bMouseRButtonDown = True
            line_start = (x, y)
            line_down_start = (x, y)

    elif (event == cv2.EVENT_MOUSEMOVE) and (flags & cv2.EVENT_FLAG_RBUTTON):
            if s_bMouseRButtonDown:
                # print("start:{0} {1} now:{2} {3}".format(line_start[0], line_start[1], x, y))
                copyImage = dstImage.copy()
                cv2.line(copyImage, (line_start[0], line_start[1]), (x, y), (255, 0, 0), 3)
                cv2.imshow(wndName, copyImage)

    elif (event == cv2.EVENT_MOUSEMOVE) and (flags & cv2.EVENT_FLAG_LBUTTON):
            cur_pt_x, cur_pt_y = x, y
            if s_bMouseLButtonDown:
                cv2.line(copyImage, (pre_pt_x,pre_pt_y), (cur_pt_x,cur_pt_y), (255, 0, 0), 3)
                cv2.imshow(wndName, copyImage)
                cv2.line(tmpImage, (pre_pt_x,pre_pt_y), (cur_pt_x,cur_pt_y), (255), 3)
                pre_pt_x,pre_pt_y = cur_pt_x,cur_pt_y

    elif event == cv2.EVENT_LBUTTONUP:
            cur_pt_x, cur_pt_y = x, y
            cv2.circle(copyImage, (cur_pt_x,cur_pt_y), 3, (0, 0, 255), -1, 0)
            dstImage = copyImage.copy()
            s_bMouseLButtonDown = False

    elif event == cv2.EVENT_RBUTTONUP:
            cv2.circle(copyImage, (x, y), 3, (0, 0, 255), -1, 0)
            cv2.line(dstImage, (line_start[0], line_start[1]), (x, y), (255, 0, 0), 3)
            cv2.line(tmpImage, (line_start[0], line_start[1]), (x, y), (255, 0, 0), 3)

            cv2.line(dstImage, (line_down_start[0], line_down_start[1]), (x, y), (255, 0, 0), 3)
            cv2.line(tmpImage, (line_down_start[0], line_down_start[1]), (x, y), (255, 0, 0), 3)

            copyImage = dstImage.copy()
            cv2.imshow(wndName, copyImage)
            s_bMouseRButtonDown = False

    elif event == cv2.EVENT_LBUTTONDBLCLK:
            tmpImage = fillHole(tmpImage)
            copyImage = draw_mask_image(srcImage,tmpImage)
            dstImage = copyImage.copy()
            cv2.imshow(wndName, tmpImage)


if __name__ == '__main__':

    file_lists = os.listdir(target_src_dir)
    global srcImage  
    for imgname in file_lists:
        srcImage = cv2.imread(target_src_dir + imgname)
        # srcImage = cv2.resize(srcImage, (512, 512), interpolation=cv2.INTER_LINEAR).copy()
        print(srcImage.shape)

        print(imgname)
        dstImage = srcImage.copy()
        copyImage = srcImage.copy()
        tmpImage = np.zeros((srcImage.shape[0], srcImage.shape[1]), np.uint8)

        cv2.namedWindow(wndName)
        cv2.setMouseCallback(wndName, draw_mask)
        key = 0
        end_flag = False
        while True:
            cv2.imshow(wndName, copyImage)  # 显示鼠标操作窗口
            key = cv2.waitKey(10)  # 读入键盘操作
            if key == ord('n'):
                break
            elif key == ord('s'):
                tmpImage = fillHole(tmpImage)
                cv2.imwrite(save_image_dir + imgname, srcImage)
                print(srcImage.shape)
                print("mask dir:{}".format(save_mask_dir + imgname, tmpImage))
                cv2.imwrite(save_mask_dir + imgname, tmpImage)
                break
            elif key == ord('d'):
                dstImage = srcImage.copy()
                copyImage = srcImage.copy()
                tmpImage = np.zeros((srcImage.shape[0], srcImage.shape[1]), np.uint8)
            elif key == 27:
                end_flag = True
                break

        if end_flag:
            break

